package prog1;

import prog1.common.Defs;
import prog1.common.Parameters;
import prog1.common.Utils;
import prog1.helpers.CipherBuilder;
import prog1.helpers.KeyStoreProxy;
import prog1.helpers.RandomIV;
import prog1.helpers.SignatureProxy;

import javax.crypto.Cipher;
import javax.crypto.CipherOutputStream;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.FileOutputStream;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.PublicKey;

/**
 * Created with IntelliJ IDEA.
 * User: Roi
 * Date: 04/12/13
 *
 * This application reads a text file (Defs.Defaults.INPUT_FILENAME), encrypts its content and digitally signs it.
 *
 * Application arguments: key store password
 *
 * Two files are created:
 * - Encrypted file (Defs.Defaults.ENCRYPTED_FILENAME)
 * - Parameters file (Defs.PARAMETERS_FILENAME)
 */
public class Encrypt {
    /**
     * Main method
     * @param args program arguments
     */
    public static void main(String[] args) {
        try {
            // Initialize parameters, use default values
            Parameters parameters = new Parameters();
            parameters.loadDefaults();

            // Read input file content
            Utils.log("Reading source file for encryption: " + parameters.get("INPUT_FILENAME"));
            String fileText = Utils.fileToString(parameters.get("INPUT_FILENAME"));

            // Retrieve key store password
            Utils.log("Reading key store password from program arguments");
            String keyStorePassword = Utils.getKeyStorePassword(args);

            // Initialize and load key store
            Utils.log("Initializing key store");
            KeyStoreProxy keyStoreProxy = new KeyStoreProxy();
            keyStoreProxy.initialize(parameters.get("KEY_STORE_TYPE"), parameters.get("KEY_STORE_PROVIDER"));
            keyStoreProxy.load(parameters.get("KEY_STORE_FILE_NAME"), keyStorePassword);

            // Generate random IV for CBC mode
            String ivStr = parameters.get("IV_SIZE_IN_BYTES");
            Utils.log("Generating random IV with size of " + ivStr + " bytes");
            RandomIV randomIV = new RandomIV(Short.parseShort(ivStr));
            randomIV.generate();

            // Generate random key
            KeyGenerator keyGenerator =
                    KeyGenerator.getInstance(parameters.get("FILE_ENCRYPTION_ALGORITHM"),
                                             parameters.get("FILE_ENCRYPTION_PROVIDER"));
            Utils.log("Generating random key");
            SecretKey randomKey = keyGenerator.generateKey();

            // Initialize cipher for encryption using CBC method and random IV
            Utils.log("Initializing cipher for file encryption");
            Cipher fileCipher = new CipherBuilder()
                    .withAlgorithm(parameters.get("FILE_ENCRYPTION_TRANSFORMATION"))
                    .withProvider(parameters.get("FILE_ENCRYPTION_PROVIDER"))
                    .withMode(Cipher.ENCRYPT_MODE)
                    .withKey(randomKey)
                    .withSpec(new IvParameterSpec(randomIV.getRandomIV()))
                    .build();

            // Retrieve private key for signature
            PrivateKey privateKey =
                    keyStoreProxy.getPrivateKey(parameters.get("KEY_STORE_SIGNATURE_ALGORITHM_ALIAS"), keyStorePassword);

            // Sign the data using given algorithm
            Utils.log("Digitally signing the data");
            SignatureProxy signatureProxy =  new SignatureProxy(parameters.get("SIGNATURE_ALGORITHM"),
                                                                parameters.get("SIGNATURE_PROVIDER"));
            signatureProxy.getSignature().initSign(privateKey);
            signatureProxy.getSignature().update(fileText.getBytes());
            signatureProxy.sign();

            // Initialize cipher output stream
            CipherOutputStream cipherOutputStream =
                    new CipherOutputStream(new FileOutputStream(parameters.get("ENCRYPTED_FILENAME")), fileCipher);

            // Write the encrypted data to file
            Utils.log("Writing encrypted file to: " + parameters.get("ENCRYPTED_FILENAME"));
            cipherOutputStream.write(fileText.getBytes());
            cipherOutputStream.close();

            // Retrieve public key
            PublicKey publicKey = keyStoreProxy.getPublicKey(parameters.get("KEY_STORE_KEY_ENCRYPTION_ALGORITHM_ALIAS"));

            // Encrypt random key with public key
            Utils.log("Initializing cipher for key encryption");
            Cipher keyCipher = new CipherBuilder()
                    .withAlgorithm(parameters.get("KEY_ENCRYPTION_ALGORITHM"))
                    .withProvider(parameters.get("KEY_ENCRYPTION_PROVIDER"))
                    .withMode(Cipher.ENCRYPT_MODE)
                    .withKey(publicKey)
                    .build();
            byte[] bytes = keyCipher.doFinal(randomKey.getEncoded());

            // Create SecretKey object from encrypted key
            SecretKey encryptedSecretKey = new SecretKeySpec(bytes, parameters.get("FILE_ENCRYPTION_ALGORITHM"));

            // Save secret key in key store
            Utils.log("Saving secret key in key store");
            keyStoreProxy.setEntry(parameters.get("FILE_ENCRYPTION_ALGORITHM"),
                    new KeyStore.SecretKeyEntry(encryptedSecretKey),
                    new KeyStore.PasswordProtection(keyStorePassword.toCharArray()));
            keyStoreProxy.save(parameters.get("KEY_STORE_FILE_NAME"), keyStorePassword);

            // Add random IV to parameters
            randomIV.serialize(parameters);

            // Add digital signature to parameters
            signatureProxy.serialize(parameters);

            // Serialize parameters to file
            Utils.log("Saving parameters to file: " + Defs.PARAMETERS_FILENAME);
            parameters.serialize();

            Utils.log("Finished");
        } catch(Exception e) {
            System.out.println("An error has occurred: " + e.getMessage());
        }
    }
}
